import os
from dotenv import load_dotenv
from cozepy import COZE_CN_BASE_URL,TokenAuth,Coze,Message,ChatEventType,COZE_CN_BASE_URL
import streamlit as st
from select import select
from langchain_core.messages import AIMessage, HumanMessage
from config_setting import chat_bot_model_config





class chatbot:


    def __init__(self):
        load_dotenv()

        coze_aip_token = os.getenv("COZE_TOKEN")
        self.coze = Coze(auth= TokenAuth(token = coze_aip_token),base_url=COZE_CN_BASE_URL)

    def get_response(self,question):
        try:

            response = self.coze.chat.stream(
                bot_id=self.model_bot_id,
                user_id="001",
                additional_messages=[
                    Message.build_user_question_text(question)
                ]
            )
            return response
        except Exception as e:
            return f'当前bot不可用，请切换其他bot继续体验'


def get_stream_response(response):
    for event in response:

        if event.event == ChatEventType.CONVERSATION_MESSAGE_DELTA:
            print(event.message.content, end="", flush=True)

            if event.message.content:
                content_str = str(event.message.content)
                yield content_str

    print('----------以上-------------')


def init_params():
    if"chat_message" not in st.session_state:
        st.session_state.chat_message = []
    if "chat_bot" not in st.session_state:
        st.session_state.chat_bot = chatbot()


def clear():
    st.session_state.chat_message=[]
    st.session_state.chat_bot = chatbot()

def chat_bot_page():
    init_params()
    with st.sidebar:
        with st.container(border=True):
            select_model = st.selectbox("选择-Bot", options=["健康生活助手"], index=0)
            model_bot_id = chat_bot_model_config.model_ls[select_model]["bot_id"]
            model_description = chat_bot_model_config.model_ls[select_model]["description"]
            st.button(label="清除聊天记录", on_click=lambda: clear(), use_container_width=True)

    st.title("健康生活助手")
    st.subheader(body='', divider='rainbow')

    '''滚动更新聊天记录'''
    with st.chat_message("BOT"):
        st.markdown(model_description)
    for message in st.session_state.chat_message:
        if isinstance(message, AIMessage):
            with st.chat_message("BOT"):
                st.markdown(message.content)
        elif isinstance(message,HumanMessage):
            with st.chat_message("Human"):
                st.markdown(message.content)

    '''用户问题交互'''
    question = st.chat_input("输入你的问题")
    if question:
        st.session_state.chat_bot.model_bot_id = model_bot_id
        with st.chat_message("Human"):
            st.markdown(question)
            st.session_state.chat_message.append(HumanMessage(content=question))
        with st.chat_message("BOT"):
            coze_res = st.session_state.chat_bot.get_response(question)
            response = st.write_stream(get_stream_response(coze_res))
            st.session_state.chat_message.append(AIMessage(content=response))






